// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/substrait/pollux_to_substrait_type.h>

#include <pollux/expression/expr.h>

namespace kumo::pollux::substrait {
    const ::substrait::Type &PolluxToSubstraitTypeConvertor::toSubstraitType(
        google::protobuf::Arena &arena,
        const pollux::TypePtr &type) {
        ::substrait::Type *substraitType =
                google::protobuf::Arena::CreateMessage<::substrait::Type>(&arena);

        switch (type->kind()) {
            case pollux::TypeKind::BOOLEAN: {
                auto substraitBool =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_Boolean>(
                            &arena);
                substraitBool->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_bool_(substraitBool);

                break;
            }
            case pollux::TypeKind::TINYINT: {
                auto substraitI8 =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_I8>(&arena);
                substraitI8->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_i8(substraitI8);
                break;
            }
            case pollux::TypeKind::SMALLINT: {
                auto substraitI16 =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_I16>(&arena);
                substraitI16->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_i16(substraitI16);
                break;
            }
            case pollux::TypeKind::INTEGER: {
                if (type->isDate()) {
                    auto substraitDate =
                            google::protobuf::Arena::CreateMessage<::substrait::Type_Date>(
                                &arena);
                    substraitDate->set_nullability(
                        ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                    substraitType->set_allocated_date(substraitDate);
                    return *substraitType;
                }

                auto substraitI32 =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_I32>(&arena);
                substraitI32->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_i32(substraitI32);
                break;
            }
            case pollux::TypeKind::BIGINT: {
                auto substraitI64 =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_I64>(&arena);
                substraitI64->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_i64(substraitI64);
                break;
            }
            case pollux::TypeKind::REAL: {
                auto substraitFp32 =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_FP32>(
                            &arena);
                substraitFp32->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_fp32(substraitFp32);
                break;
            }
            case pollux::TypeKind::DOUBLE: {
                auto substraitFp64 =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_FP64>(
                            &arena);
                substraitFp64->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_fp64(substraitFp64);
                break;
            }
            case pollux::TypeKind::VARCHAR: {
                auto substraitString =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_String>(
                            &arena);
                substraitString->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_string(substraitString);
                break;
            }
            case pollux::TypeKind::VARBINARY: {
                auto substraitVarBinary =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_Binary>(
                            &arena);
                substraitVarBinary->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_binary(substraitVarBinary);
                break;
            }
            case pollux::TypeKind::TIMESTAMP: {
                auto substraitTimestamp =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_Timestamp>(
                            &arena);
                substraitTimestamp->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_timestamp(substraitTimestamp);
                break;
            }
            case pollux::TypeKind::ARRAY: {
                ::substrait::Type_List *substraitList =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_List>(
                            &arena);

                substraitList->mutable_type()->MergeFrom(
                    toSubstraitType(arena, type->as_array().elementType()));

                substraitList->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);

                substraitType->set_allocated_list(substraitList);

                break;
            }
            case pollux::TypeKind::MAP: {
                ::substrait::Type_Map *substraitMap =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_Map>(&arena);

                substraitMap->mutable_key()->MergeFrom(
                    toSubstraitType(arena, type->as_map().keyType()));
                substraitMap->mutable_value()->MergeFrom(
                    toSubstraitType(arena, type->as_map().valueType()));

                substraitMap->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);

                substraitType->set_allocated_map(substraitMap);

                break;
            }
            case pollux::TypeKind::ROW: {
                ::substrait::Type_Struct *substraitStruct =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_Struct>(
                            &arena);
                for (const auto &child: type->as_row().children()) {
                    substraitStruct->set_nullability(
                        ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                    substraitStruct->add_types()->MergeFrom(toSubstraitType(arena, child));
                }
                substraitStruct->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_struct_(substraitStruct);
                break;
            }
            case pollux::TypeKind::UNKNOWN: {
                auto substraitUserDefined =
                        google::protobuf::Arena::CreateMessage<::substrait::Type_UserDefined>(
                            &arena);
                substraitUserDefined->set_type_reference(0);
                substraitUserDefined->set_nullability(
                    ::substrait::Type_Nullability_NULLABILITY_NULLABLE);
                substraitType->set_allocated_user_defined(substraitUserDefined);
                break;
            }
            case pollux::TypeKind::FUNCTION:
            case pollux::TypeKind::OPAQUE:
            case pollux::TypeKind::INVALID:
            default:
                POLLUX_UNSUPPORTED("Unsupported pollux type '{}'", type->toString());
        }
        return *substraitType;
    }

    const ::substrait::NamedStruct &
    PolluxToSubstraitTypeConvertor::toSubstraitNamedStruct(
        google::protobuf::Arena &arena,
        const pollux::RowTypePtr &rowType) {
        ::substrait::NamedStruct *substraitNamedStruct =
                google::protobuf::Arena::CreateMessage<::substrait::NamedStruct>(&arena);

        const auto size = rowType->size();
        if (size != 0) {
            const auto &names = rowType->names();
            const auto &polluxTypes = rowType->children();

            auto substraitType = substraitNamedStruct->mutable_struct_();

            substraitType->set_nullability(
                ::substrait::Type_Nullability_NULLABILITY_NULLABLE);

            for (int64_t i = 0; i < size; ++i) {
                const auto &name = names.at(i);
                const auto &polluxType = polluxTypes.at(i);
                substraitNamedStruct->add_names(name);

                substraitType->add_types()->MergeFrom(toSubstraitType(arena, polluxType));
            }
        }

        return *substraitNamedStruct;
    }
} // namespace kumo::pollux::substrait
